from .config_manager import ConfigManager
# from asyncio import tasks
import os
"""
命令执行器模块

负责任务调度和命令执行，支持串行、并行执行模式
处理任务依赖关系，提供实时执行状态反馈
"""

import asyncio
import logging
from colorama import Fore, Style, init
import paramiko
from typing import List, Dict, Any
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

import traceback

from .config_manager import Task, Host
from .ssh_connection import SSHConnectionManager
from .result_collector import TaskResult, CommandResult
from .stream_executor import streamCommandExecutor, StreamingResult
from .command_classifier import CommandClassifier, CommandType
from .expected_match import ExpectedMatch

#
init(autoreset=True)

#  获取当前模块的日志记录器
logger = logging.getLogger(__name__)

# logging.basicConfig(level=logging.INFO)

class CommandExecutor:
    """
    命令执行器
    
    功能：
    1. 管理任务的执行流程和调度
    2. 支持串行和并行两种执行模式
    3. 处理任务间的依赖关系
    4. 提供异步执行能力以提高效率
    5. 收集和整理执行结果
    6. 异常处理和错误恢复
    """
    
    def __init__(self, ssh_manager: SSHConnectionManager, max_workers: int = 10):
        """
        初始化命令执行器
        
        Args:
            ssh_manager: SSH连接管理器实例
            max_workers: 最大工作线程数，用于并发执行
        """
        self.ssh_manager = ssh_manager           # SSH连接管理器
        self.max_workers = max_workers           # 最大并发数
        self.executor = ThreadPoolExecutor(max_workers=max_workers)  # 线程池
        # self.stream_executor = streamCommandExecutor(ssh_manager)
        self.ssh_manager_stream = streamCommandExecutor(ssh_manager)
        self.classifier = CommandClassifier()
        
        # 任务执行结果缓存，用于依赖检查
        # key: 任务名称, value: TaskResult对象
        self.task_results: Dict[str, TaskResult] = {}
        
        # 执行统计
        self.total_tasks = 0
        self.completed_tasks = 0
        self.failed_tasks = 0
        self.success_tasks = 0
    

    async def execute_tasks(self, config_manager: ConfigManager, progress_callback=None,hosts_input:Dict[str, Host]={}) -> List[TaskResult]:
        """
        执行所有任务的主入口方法
        
        按照任务依赖关系和执行模式，协调所有任务的执行
        
        Args:
            tasks: 要执行的任务列表
            hosts: 主机配置字典
            progress_callback: 进度回调函数 (completed_count, total_count)
            
        Returns:
            List[TaskResult]: 所有任务的执行结果
        """
        tasks = config_manager.tasks
        # hosts = config_manager.hosts
        if config_manager.no_host:
            if hosts_input:
                hosts = hosts_input
            else:
                logger.error("execute_tasks|未提供主机配置，且配置文件中未定义主机")
                return []
        else:
            hosts = config_manager.hosts
        self.config_manager = config_manager

        self.total_tasks = len(tasks)
        results = []
        
        logger.info(f"开始执行 {self.total_tasks} 个任务")
        
        # 初始进度回调
        if progress_callback:
            logger.info(f"execute_tasks|初始进度: 0/{self.total_tasks}")
            progress_callback(0, self.total_tasks)
        else:
            logger.info(f"execute_tasks|无进度回调函数")

        # 按顺序执行每个任务
        for i, task in enumerate(tasks, 1):
            logger.info(f"[{i}/{self.total_tasks}] 准备执行任务: {task.name}")
            
            # 检查任务依赖是否满足
            if not self._check_dependencies(task):
                logger.warning(f"任务 {task.name} 的依赖项执行失败，跳过执行")
                
                # 创建一个跳过状态的结果
                skipped_result = self._create_empty_result(task)
                self.task_results[task.name] = skipped_result
                results.append(skipped_result)
                continue
            
            # 获取任务需要执行的主机列表
            if config_manager.no_host:
                task_hosts = list(hosts_input.values())
            else:
                task_hosts = self._get_valid_hosts(task, hosts)
            
            if not task_hosts:
                logger.error(f"任务 {task.name} 没有可用的主机，跳过执行")
                
                # 创建一个无主机的失败结果
                no_host_result = self._create_empty_result(task)
                self.task_results[task.name] = no_host_result
                results.append(no_host_result)
                continue
            
            # 根据执行模式选择执行方法
            try:
                if task.mode == "parallel":
                    task_result = await self._execute_parallel(task, task_hosts)
                else:  # serial 模式
                    task_result = await self._execute_serial(task, task_hosts)
                
                # 缓存结果用于依赖检查
                self.task_results[task.name] = task_result
                results.append(task_result)
                
                # 记录任务完成状态
                status = "成功" if task_result.success else "失败"
                logger.info(f"任务 {task.name} 执行{status}")

                if task_result.success:
                    # 任务成功，执行后续任务
                    self.success_tasks += 1
                    continue
                else:
                    # 任务失败，根据配置决定是否继续执行后续任务
                    if task.allow_failure:
                        self.failed_tasks += 1
                        continue
                    else:
                        # break
                        pass
                
            except Exception as e:
                logger.error(f"任务 {task.name} 执行异常: {str(e)}")
                
                # 创建异常结果
                exception_result = self._create_exception_result(task, str(e))
                self.task_results[task.name] = exception_result
                results.append(exception_result)
                self.failed_tasks += 1
            
            # 进度回调
            if progress_callback: # main.py中构建的
                progress_callback(i, self.total_tasks)
        
        # 记录总体执行统计
        logger.info(f"任务执行完成: 总数={self.total_tasks}, "
                   f"成功={self.success_tasks}, "
                   f"失败={self.failed_tasks}, "
                   f"跳过={self.total_tasks-self.success_tasks-self.failed_tasks}")
        
        return results
    # keep this function
    def _check_dependencies(self, task: Task) -> bool:
        """
        检查任务的依赖关系是否满足
        
        Args:
            task: 要检查的任务
            
        Returns:
            bool: 依赖是否满足
        """
        # 如果没有依赖，直接返回True
        if not task.depends_on:
            return True
        
        logger.debug(f"检查任务 {task.name} 的依赖: {task.depends_on}")
        
        for dep_name in task.depends_on:
            # 检查依赖任务是否已执行
            if dep_name not in self.task_results:
                logger.warning(f"依赖任务 {dep_name} 尚未执行")
                return False
            
            dep_result = self.task_results[dep_name]
            
            # 如果依赖任务失败且当前任务不允许失败，则依赖不满足
            if not dep_result.success and not task.allow_failure:
                logger.warning(f"依赖任务 {dep_name} 执行失败，且任务 {task.name} 不允许失败")
                return False
        
        logger.debug(f"任务 {task.name} 的所有依赖都已满足")
        return True
    
    # keep this function
    def _get_valid_hosts(self, task: Task, hosts: Dict[str, Host]) -> List[Host]:
        """
        获取任务的有效主机列表
        
        Args:
            task: 任务对象
            hosts: 主机配置字典
            
        Returns:
            List[Host]: 有效的主机列表
        """
        valid_hosts = []
        
        for host_name in task.hosts:
            if host_name in hosts:
                valid_hosts.append(hosts[host_name])
            else:
                logger.warning(f"任务 {task.name} 指定的主机不存在: {host_name}")
        
        return valid_hosts
    # keep this function
    async def _execute_serial(self, task: Task, hosts: List[Host]) -> TaskResult:
        """
        串行执行任务（按主机顺序依次执行）
        
        Args:
            task: 任务配置
            hosts: 主机列表
            
        Returns:
            TaskResult: 任务执行结果
        """
        logger.info(f"开始串行执行任务: {task.name} (主机数: {len(hosts)})")
        
        # 创建任务结果对象
        task_result = TaskResult(
            task_name=task.name,
            start_time=datetime.now()
        )
        
        # 逐个主机执行命令
        for i, host in enumerate(hosts, 1):
            logger.info(f"[{i}/{len(hosts)}] 在主机 {host.name} 上执行命令")
            
            # 在单个主机上执行命令
            cmd_result = await self._execute_on_host(task, host)
            task_result.command_results.append(cmd_result)
            
            # 如果命令失败且不允许失败，停止后续执行
            if not cmd_result.success and not task.allow_failure:
                logger.error(f"任务 {task.name} 在主机 {host.name} 上执行失败，停止后续主机的执行")
                break
        
        # 设置任务完成时间和成功状态
        task_result.end_time = datetime.now()
        task_result.success = self._calculate_task_success(task_result.command_results, task.allow_failure)
        
        duration = (task_result.end_time - task_result.start_time).total_seconds()
        logger.info(f"任务 {task.name} 串行执行完成，耗时 {duration:.2f}秒")
        
        return task_result
    # keep this function
    async def _execute_parallel(self, task: Task, hosts: List[Host]) -> TaskResult:
        """
        并行执行任务（同时在所有主机上执行）
        
        Args:
            task: 任务配置
            hosts: 主机列表
            
        Returns:
            TaskResult: 任务执行结果
        """
        logger.info(f"开始并行执行任务: {task.name} (主机数: {len(hosts)})")
        
        # 创建任务结果对象
        task_result = TaskResult(
            task_name=task.name,
            start_time=datetime.now()
        )
        
        # 创建并行任务列表
        parallel_tasks = [self._execute_on_host(task, host) for host in hosts]
        
        try:
            # 并行执行所有任务，收集结果
            results = await asyncio.gather(*parallel_tasks, return_exceptions=True)
            
            # 处理执行结果
            for i, result in enumerate(results):
                if isinstance(result, Exception):

                    # 处理执行过程中的异常
                    logger.error(f"主机 {hosts[i].name} 执行异常: {result}")
                    logger.error(traceback.format_exc())
                    
                    # 创建异常命令结果
                    exception_cmd_result = CommandResult(
                        host_name=hosts[i].name,
                        command=task.command,
                        exit_code=-1,
                        stdout="",
                        stderr=f"工具运行异常: {str(result)}",
                        success=False,
                        start_time=datetime.now(),
                        end_time=datetime.now()
                    )
                    task_result.command_results.append(exception_cmd_result)
                    raise result
                elif isinstance(result, CommandResult):  # 添加类型检查
                    # 正常的命令结果
                    task_result.command_results.append(result)
                else:
                    # 处理意外的结果类型
                    logger.error(f"意外的结果类型: {type(result)}")
        
        except Exception as e:
            logger.error(f"并行执行任务 {task.name} 时发生异常: {str(e)}")
            logger.error(traceback.format_exc())
        
        # 设置任务完成时间和成功状态
        task_result.end_time = datetime.now()
        task_result.success = self._calculate_task_success(task_result.command_results, task.allow_failure)
        
        duration = (task_result.end_time - task_result.start_time).total_seconds()
        logger.info(f"任务 {task.name} 并行执行完成，耗时 {duration:.2f}秒")
        
        return task_result
    # keep this function
    async def _execute_on_host(self, task: Task, host: Host) -> CommandResult:
        """
        在单个主机上执行命令
        
        Args:
            task: 任务配置
            host: 主机配置
            
        Returns:
            CommandResult: 命令执行结果
        """
        # 分类命令类型
        cmd_type_input = task.command_type
        cmd_type = self.classifier.classify_command(task.command,cmd_type_input)


        logger.info(f"{task.name} 在主机 {host.name} 上执行命令，类型: {cmd_type}")

        if cmd_type == CommandType.STREAMING:
            return await self._execute_streaming_command(task, host)
        elif cmd_type == CommandType.LONG_RUNNING:
            return await self._execute_long_running_command(task, host)
        else:
            return await self._execute_normal_command(task, host)

    
    # keep this function
    async def _execute_normal_command(self, task: Task, host: Host) -> CommandResult:
        """
        执行普通命令（原有逻辑的增强版本）
        
        Args:
            task: 任务配置
            host: 主机配置
            
        Returns:
            CommandResult: 命令执行结果
        """
        start_time = datetime.now()
        logger.debug(f"执行普通命令 [{host.name}]: {task.command}")
        
        try:
            # 在线程池中执行阻塞的SSH操作，避免阻塞事件循环
            exit_code, stdout, stderr = await asyncio.get_event_loop().run_in_executor(
                self.executor,
                self.ssh_manager.execute_command,
                host,
                task.command,
                task.timeout
            )
    
            
            end_time = datetime.now()
            
            # 判断命令是否成功执行
            success = exit_code == 0
            
            # 如果命令成功但有期望输出要求，检查输出内容
            if success and (task.expected or task.unexpected):
                expected_match = ExpectedMatch(task.expected, task.unexpected)
                expected_match.match(stdout+stderr)
                ok,exception = expected_match.check()
                if not ok:
                    success = False
                    stderr += f"\n预期输出验证失败: '{exception}'"
                    logger.warning(f"命令执行成功但输出验证失败 [{host.name}|{task.name}] 异常: '{exception}'")

  
            
            # 创建命令结果对象
            cmd_result = CommandResult(
                host_name=host.name,
                command=task.command,
                exit_code=exit_code,
                stdout=stdout,
                stderr=stderr,
                success=success,
                start_time=start_time,
                end_time=end_time
            )
            
            # 记录执行结果
            duration = (end_time - start_time).total_seconds()
            if success:
                logger.debug(f"普通命令执行成功 [{host.name}] 耗时{duration:.2f}s: {task.command}")
            else:
                logger.warning(f"普通命令执行失败 [{host.name}] 退出码{exit_code}: {task.command}")
                if stderr:
                    logger.debug(f"错误输出 [{host.name}]: {stderr}")
            
            return cmd_result
            
        except Exception as e:
            end_time = datetime.now()
            error_msg = str(e)
            logger.error(f"普通命令执行异常 [{host.name}]: {error_msg}")
            
            # 创建异常结果
            return CommandResult(
                host_name=host.name,
                command=task.command,
                exit_code=-1,
                stdout="",
                stderr=error_msg,
                success=False,
                start_time=start_time,
                end_time=end_time
            )
        

    async def _execute_long_running_command(self, task: Task, host: Host) -> CommandResult:
        """
        执行长时间运行的命令（原有逻辑的增强版本）
        
        Args:
            task: 任务配置
            host: 主机配置
            
        Returns:
            CommandResult: 命令执行结果
        """
        start_time = datetime.now()
        logger.debug(f"执行长时间运行的命令 [{host.name}]: {task.command}")

        # 创建长时间运行命令的进度监控任务
        progress_task = asyncio.create_task(
                self._monitor_long_running_progress(host.name, task.name, task.timeout)
        )
        
        try:
            # 在线程池中执行阻塞的SSH操作，避免阻塞事件循环
            exit_code, stdout, stderr = await asyncio.get_event_loop().run_in_executor(
                self.executor,
                self.ssh_manager.execute_command,
                host,
                task.command,
                task.timeout
            )

            end_time = datetime.now()
            
            # 判断命令是否成功执行
            success = exit_code == 0
            
            # 如果命令成功但有期望输出要求，检查输出内容
            if success and (task.expected or task.unexpected):
                expected_match = ExpectedMatch(task.expected, task.unexpected)
                expected_match.match(stdout+stderr)
                ok,exception = expected_match.check()
                if not ok:
                    success = False
                    stderr += f"\n预期输出验证失败: '{exception}'"
                    logger.warning(f"长时间运行的命令执行成功但输出验证失败 [{host.name}|{task.name}] 异常: '{exception}'")


            # 取消进度监控任务
            progress_task.cancel()
            
            # 创建命令结果对象
            cmd_result = CommandResult(
                host_name=host.name,
                command=task.command,
                exit_code=exit_code,
                stdout=stdout,
                stderr=stderr,
                success=success,
                start_time=start_time,
                end_time=end_time
            )
            
            # 记录执行结果
            duration = (end_time - start_time).total_seconds()
            if success:
                logger.debug(f"长时间运行的命令执行成功 [{host.name}] 耗时{duration:.2f}s: {task.command}")
            else:
                logger.warning(f"长时间运行的命令执行失败 [{host.name}] 退出码{exit_code}: {task.command}")
                if stderr:
                    logger.debug(f"错误输出 [{host.name}]: {stderr}")
            
            return cmd_result
            
        except Exception as e:
            end_time = datetime.now()
            error_msg = str(e)
            logger.error(f"长时间运行的命令执行异常 [{host.name}]: {error_msg}")
            
            # 取消进度监控任务
            progress_task.cancel()
            
            # 创建异常结果
            return CommandResult(
                host_name=host.name,
                command=task.command,
                exit_code=-1,
                stdout="",
                stderr=error_msg,
                success=False,
                start_time=start_time,
                end_time=end_time
            )
    
    async def _execute_streaming_command(self, task: Task, host: Host) -> CommandResult:
        """执行流式命令"""
        start_time = datetime.now()
        
        # 获取流式配置
        # cmd_output_dir = self.config_manager.cmd_output_dir

        # streaming_config = getattr(task, 'streaming_config', {})
        # max_duration = streaming_config.get('max_duration', 300)
        # show_realtime = streaming_config.get('show_realtime', False)
        # write_to_file = streaming_config.get('write_to_file', True)


        # auto_stop = streaming_config.get('auto_stop', False) # todo
        
        # 实时输出回调
        output_callback = None
        # if show_realtime:
        #     output_callback = self._create_realtime_callback(task.name, host.name)
        
      
        # 执行流式命令
        streaming_result = await self.ssh_manager_stream.execute_streaming_command(
            host=host,
            task=task,
            cmd_output_dir= self.config_manager.cmd_output_dir
        )
        
        # 将StreamingResult转换为标准CommandResult
        end_time = datetime.now()
        # 查看文件是否存在
        if streaming_result.stdout_file and os.path.exists(streaming_result.stdout_file):
            stdout = "命令流式输出内容重定向到: "+streaming_result.stdout_file +"\n"+ '\n'.join(streaming_result.output_lines)
        else:
            stdout = "未能将流式输出内容写入到文件"+"\n"+ '\n'.join(streaming_result.output_lines) 

        if streaming_result.stderr_file and os.path.exists(streaming_result.stderr_file):
            stderr = "命令流式错误内容重定向到: "+streaming_result.stderr_file +"\n"+ '\n'.join(streaming_result.error_lines)
        else:
            stderr = "未能将流式错误内容写入到文件"+"\n"+ '\n'.join(streaming_result.error_lines)
        stderr = '\n'.join(streaming_result.error_lines)

        return CommandResult(
            host_name=host.name,
            command=task.command,
            exit_code=0 if not streaming_result.error_lines else 1,
            stdout=stdout,
            stderr=stderr,
            success=len(streaming_result.error_lines) == 0,
            start_time=start_time,
            end_time=end_time
        )
        
    
    
    def _calculate_task_success(self, command_results: List[CommandResult], allow_failure: bool) -> bool:
        """
        计算任务的整体成功状态
        
        Args:
            command_results: 命令执行结果列表
            allow_failure: 是否允许失败
            
        Returns:
            bool: 任务是否成功
        """
        if not command_results:
            return False
        
        if allow_failure:
            # 如果允许失败，只要有一个成功就算成功
            return any(result.success for result in command_results)
        else:
            # 如果不允许失败，所有命令都必须成功
            return all(result.success for result in command_results)
    
   
    
    def _create_empty_result(self, task: Task) -> TaskResult:
        """
        创建空的任务结果
        
        Args:
            task: 任务配置
            
        Returns:
            TaskResult: 空的任务结果
        """
        now = datetime.now()
        return TaskResult(
            task_name=task.name,
            start_time=now,
            end_time=now,
            success=False
        )

    @staticmethod
    def create_none_result() -> TaskResult:
        """
        创建空的任务结果
        
        Args:
            task: 任务配置
            
        Returns:
            TaskResult: 空的任务结果
        """
        now = datetime.now()
        return TaskResult(
            task_name='',
            start_time=now,
            end_time=now,
            success=False
        )
    
    def _create_exception_result(self, task: Task, error_msg: str) -> TaskResult:
        """
        创建异常状态的任务结果
        
        Args:
            task: 任务配置
            error_msg: 错误信息
            
        Returns:
            TaskResult: 异常状态的任务结果
        """
        now = datetime.now()
        result = TaskResult(
            task_name=task.name,
            start_time=now,
            end_time=now,
            success=False
        )
        
        # 添加一个表示异常的命令结果
        exception_cmd = CommandResult(
            host_name="unknown",
            command=task.command,
            exit_code=-1,
            stdout="",
            stderr=error_msg,
            success=False,
            start_time=now,
            end_time=now
        )
        result.command_results.append(exception_cmd)
        
        return result
    
    
    async def _monitor_long_running_progress(
        self, 
        host_name: str, 
        task_name: str, 
        total_timeout: int
    ):
        """
        监控长时间运行命令的进度
        
        Args:
            host_name: 主机名
            task_name: 任务名
            total_timeout: 总超时时间
        """
        start_time = asyncio.get_event_loop().time()
        report_interval = min(30, total_timeout // 10)  # 进度报告间隔

        try:
            while True:
                await asyncio.sleep(report_interval)
                
                current_time = asyncio.get_event_loop().time()
                elapsed = current_time - start_time
                remaining = total_timeout - elapsed
                
                if remaining <= 0:
                    break
                
                # 计算进度百分比
                progress = (elapsed / total_timeout) * 100
                
                logger.info(
                    f"长时间任务进度 [{host_name}] {task_name}: "
                    f"{progress:.1f}% (已运行{elapsed:.0f}s, 剩余允许{remaining:.0f}s)"
                )
                
        except asyncio.CancelledError:
            logger.debug(f"进度监控已取消 [{host_name}] {task_name}")
    

    
    def _create_realtime_callback(self, task_name: str, host_name: str):
        """
        为流式命令创建实时输出回调函数
        
        Args:
            task_name: 任务名称
            host_name: 主机名称
            
        Returns:
            function: 回调函数
        """
        async def callback(line: str, result):
            """实时输出回调"""
            # 过滤掉过长的行，避免显示混乱
            display_line = line[:100] + "..." if len(line) > 100 else line
            
            # 使用不同颜色区分不同主机的输出
            colors = ["blue", "green", "yellow", "magenta", "cyan"]
            color_index = hash(host_name) % len(colors)
            color = colors[color_index]
            
            # 实时显示（这里可以集成到Rich Live显示中）
            logger.info(f"[{color}]{host_name}[/{color}] | {display_line}")
        

            # 实时显示，使用 colorama 输出颜色
            logger.info(f"[{host_name}] | {display_line}")

        return callback
    

    
    def cancel_all_tasks(self):
        """
        取消所有正在运行的长时间任务
        
        用于程序关闭时的清理
        """
        logger.info("正在取消所有任务...")
        
        # 通知ssh_manager_stream停止所有流式命令
        self.ssh_manager_stream.stop_all_streaming()
    
    def get_task_statistics(self) -> dict:
        """
        获取任务执行统计，包括不同类型命令的统计
        
        Returns:
            dict: 统计信息
        """
        base_stats = self.get_execution_stats()
        
        # 可以添加更详细的分类统计
        return {
            **base_stats,
            'command_type_stats': {
                'normal_commands': 0,      # 这些需要在执行过程中统计
                'streaming_commands': 0,
                'long_running_commands': 0
            }
        }
    
    def get_execution_stats(self) -> Dict[str, Any]:
        """
        获取执行统计信息
        
        Returns:
            Dict[str, Any]: 执行统计数据
        """
        return {
            'total_tasks': self.total_tasks,
            'success_tasks': self.success_tasks,
            'failed_tasks': self.failed_tasks,
            'success_rate': f"{((self.success_tasks) / max(self.total_tasks, 1) * 100):.1f}%",
            'cached_results': len(self.task_results)
        }
    
    def __del__(self):
        """
        析构函数：清理线程池资源
        """
        try:
            # 关闭线程池，等待所有任务完成
            self.executor.shutdown(wait=True)
        except:
            pass  # 析构时忽略异常

